#pragma once
#include <fstream>
#include <string>

using namespace std;

class FileStreamIn : public ifstream
{
public:
  FileStreamIn (void) : ifstream () {}
  explicit FileStreamIn (const string& filename, 
    ios_base::openmode mode = in | binary)
    : ifstream (filename.c_str (), mode) {}

  // Read binary data
  template< class _Ty >
  FileStreamIn& bin_read (_Ty* data, size_t count = 1)
  {
    read ((char*)data, count * sizeof (_Ty));
    return *this;
  }

  // Read string
  FileStreamIn& str_read (string& str, size_t length)
  {
    char* temp = new char[length + 1];
    read (temp, length);
    temp[length] = '\0';
    str = temp;
    delete[] temp;
    return *this;
  }

  bool check_file_type (string file_type)
  {
    seekg (0);
    string file_type_read;
    str_read (file_type_read, file_type.size ());
    if (file_type_read == file_type)
      return true;
    else
      return false;
  }
};

class FileStreamOut : public ofstream
{
public:
  FileStreamOut (void) : ofstream () {}
  explicit FileStreamOut (const string& filename, 
    ios_base::openmode mode = out | binary )
    : ofstream (filename.c_str (), mode) {}

  // Write binary data
  template< class _Ty >
  FileStreamOut& bin_write (const _Ty* data, size_t count = 1)
  {
    write ((const char*)data, count * sizeof (_Ty));
    return *this;
  }

  // Write string
  FileStreamOut& str_write (const string& str, size_t length)
  {
    if (str.size () + 1 <= length)
    {
      write (str.c_str (), str.size () + 1);
      char ch = '\0';
      for (size_t i = str.size () + 1; i < length; ++i)
        write (&ch, 1);
    }
    else
      write (str.c_str (), length);
    return *this;
  }
};

class FileStream : public fstream
{
public:
  FileStream (void) : fstream () {}
  explicit FileStream (const string& filename, 
    ios_base::openmode mode = in | out | binary )
    : fstream (filename.c_str (), mode) {}

  // Write binary data
  template< class _Ty >
  FileStream& bin_write (const _Ty* data, size_t count = 1)
  {
    write ((const char*)data, count * sizeof (_Ty));
    return *this;
  }

  // Write string
  FileStream& str_write (const string& str, size_t length)
  {
    if (str.size () + 1 <= length)
    {
      write (str.c_str (), str.size () + 1);
      char ch = '\0';
      for (size_t i = str.size () + 1; i < length; ++i)
        write (&ch, 1);
    }
    else
      write (str.c_str (), length);
    return *this;
  }

  // Read binary data
  template< class _Ty >
  FileStream& bin_read (_Ty* data, size_t count = 1)
  {
    read ((char*)data, count * sizeof (_Ty));
    return *this;
  }

  // Read string
  FileStream& str_read (string& str, size_t length)
  {
    char* temp = new char[length + 1];
    read (temp, length);
    temp[length] = '\0';
    str = temp;
    delete[] temp;
    return *this;
  }

  bool check_file_type (string file_type)
  {
    seekg (0);
    string file_type_read;
    str_read (file_type_read, file_type.size ());
    if (file_type_read == file_type)
      return true;
    else
      return false;
  }
};